#!/usr/bin/env python
# coding: utf-8

# In[14]:


import datetime
import requests

class EtigerDataClient:
    def __init__(self, base_url):
        self.__base_url = base_url
        self.__szse_calendar_url = self.__base_url + '/finance/szse/calendar'
        self.__szse_trade_date_url = self.__base_url + '/finance/szse/trade_date'
        self.__stock_kline_query_url = self.__base_url + '/finance/stock/kline/query'
        self.__stock_dividend_url = self.__base_url + '/finance/stock/dividend'
        self.__stock_lhb_url = self.__base_url + '/finance/stock/lhb'
        self.__stock_dzjy_url = self.__base_url + '/finance/stock/dzjy'
        self.__stock_bulletin_url = self.__base_url + '/finance/stock/bulletin'
    def __parse_json_response(self, rsp):
        code = rsp['code']
        msg = rsp['msg']
        trace_id = rsp['traceId']
        if '0' == code:
            return rsp['data']
        else:
            print("[%s][ERROR] EtigerDataClient response error, trace_id [%s]" %(datetime.datetime.now(), trace_id))
            raise Exception(msg)
    def get_szse_calendar(self, date_begin=None, date_end=None, open_close=None, asc=None):
        '''
        查询深交所交易日历
        '''
        payload = {
            'dateBegin': date_begin,
            'dateEnd': date_end,
            'open': open_close,
            'asc': asc
        }
        rsp = requests.get(self.__szse_calendar_url, params=payload, verify=True)
        rsp = rsp.json()
        return self.__parse_json_response(rsp)
    def get_szse_trade_date(self, date_begin=None, date_end=None, asc=None):
        '''
        查询深交所交易日
        '''
        payload = {
            'dateBegin': date_begin,
            'dateEnd': date_end,
            'asc': asc
        }
        rsp = requests.get(self.__szse_trade_date_url, params=payload, verify=True)
        rsp = rsp.json()
        return self.__parse_json_response(rsp)
    def stock_kline_query(self, batch_list=None, key_list=None):
        '''
        查询A股K线数据
        '''
        payload = {
            'batchList': batch_list,
            'keyList': key_list
        }
        rsp = requests.post(self.__stock_kline_query_url, json=payload, verify=True)
        rsp = rsp.json()
        return self.__parse_json_response(rsp)
    def get_stock_dividend(self, instrument_id=None, register_date_begin=None, register_date_end=None, dr_date_begin=None, dr_date_end=None):
        '''
        查询上市公司送股分红数据
        '''
        payload = {
            'instrumentId': instrument_id,
            'registerDateBegin': register_date_begin,
            'registerDateEnd': register_date_end,
            'drDateBegin': dr_date_begin,
            'drDateEnd': dr_date_end
        }
        rsp = requests.get(self.__stock_dividend_url, params=payload, verify=True)
        rsp = rsp.json()
        return self.__parse_json_response(rsp)
    def get_stock_lhb(self, instrument_id=None, date_begin=None, date_end=None):
        '''
        查询股票龙虎榜数据
        '''
        payload = {
            'instrumentId': instrument_id,
            'dateBegin': date_begin,
            'dateEnd': date_end
        }
        rsp = requests.get(self.__stock_lhb_url, params=payload, verify=True)
        rsp = rsp.json()
        return self.__parse_json_response(rsp)
    def get_stock_lhb_by_date(self, date):
        '''
        获取指定日期的股票龙虎榜数据
        '''
        url = self.__stock_lhb_url + "/" + date
        rsp = requests.get(url, verify=True)
        rsp = rsp.json()
        return self.__parse_json_response(rsp)
    def get_stock_dzjy(self, instrument_id=None, date_begin=None, date_end=None):
        '''
        查询股票大宗交易数据
        '''
        payload = {
            'instrumentId': instrument_id,
            'dateBegin': date_begin,
            'dateEnd': date_end
        }
        rsp = requests.get(self.__stock_dzjy_url, params=payload, verify=True)
        rsp = rsp.json()
        return self.__parse_json_response(rsp)
    def get_stock_dzjy_by_date(self, date):
        '''
        获取指定日期的股票大宗交易数据
        '''
        url = self.__stock_dzjy_url + "/" + date
        rsp = requests.get(url, verify=True)
        rsp = rsp.json()
        return self.__parse_json_response(rsp)
    def get_stock_bulletin(self, instrument_id=None, date_begin=None, date_end=None):
        '''
        查询上市公司公告数据
        '''
        payload = {
            'instrumentId': instrument_id,
            'dateBegin': date_begin,
            'dateEnd': date_end
        }
        rsp = requests.get(self.__stock_bulletin_url, params=payload, verify=True)
        rsp = rsp.json()
        return self.__parse_json_response(rsp)

def example_get_szse_calendar(data_client):
    result = data_client.get_szse_calendar('2010-01-01', '2010-01-31', 0, True)
    print("[%s][INFO] get_szse_calendar [%s]" %(datetime.datetime.now(), result))

def example_get_szse_trade_date(data_client):
    result = data_client.get_szse_trade_date('2010-01-01', '2010-01-31', False)
    print("[%s][INFO] get_szse_trade_date [%s]" %(datetime.datetime.now(), result))

def example_stock_kline_query(data_client):
    batch_list = [
        {
            "instrumentIdList":[
                "601318.SSE",
                "000001.SZSE",
                "600036.SSE"
            ],
            "tradingDateBegin":"2021-06-01",
            "tradingDateEnd":"2021-06-10"
        },
        {
            "instrumentIdList":[
                "600143.SSE"
            ],
            "tradingDateBegin":"2021-05-01",
            "tradingDateEnd":"2021-05-10"
        }
    ]
    key_list = [
        {
            "tradingDate":"2021-06-17",
            "instrumentId":"601318.SSE"
        },
        {
            "tradingDate":"2021-06-18",
            "instrumentId":"000001.SZSE"
        }
    ]
    result = data_client.stock_kline_query(batch_list, key_list)
    print("[%s][INFO] stock_kline_query [%s]" %(datetime.datetime.now(), result))

def example_get_stock_dividend(data_client):
    result = data_client.get_stock_dividend('000001.SZSE', '2010-01-01', '2021-12-31', None, None)
    print("[%s][INFO] get_stock_dividend [%s]" %(datetime.datetime.now(), result))

def example_get_stock_lhb(data_client):
    result = data_client.get_stock_lhb('000001.SZSE', '2017-06-14', '2017-08-14')
    print("[%s][INFO] get_stock_lhb [%s]" %(datetime.datetime.now(), result))

def example_get_stock_lhb_by_date(data_client):
    result = data_client.get_stock_lhb_by_date('2021-08-06')
    print("[%s][INFO] get_stock_lhb_by_date [%s]" %(datetime.datetime.now(), result))

def example_get_stock_dzjy(data_client):
    result = data_client.get_stock_dzjy('000001.SZSE', '2021-06-01', '2021-08-31')
    print("[%s][INFO] get_stock_dzjy [%s]" %(datetime.datetime.now(), result))

def example_get_stock_dzjy_by_date(data_client):
    result = data_client.get_stock_dzjy_by_date('2021-08-06')
    print("[%s][INFO] get_stock_dzjy_by_date [%s]" %(datetime.datetime.now(), result))

def example_get_stock_bulletin(data_client):
    result = data_client.get_stock_bulletin('000001.SZSE', '2021-06-01', '2021-08-31')
    print("[%s][INFO] get_stock_bulletin [%s]" %(datetime.datetime.now(), result))

if __name__ == '__main__':
    base_url = 'https://ssl.etigerfund.cn/etiger/data/api/'
    data_client = EtigerDataClient(base_url)
    #example_get_szse_calendar(data_client)
    #example_get_szse_trade_date(data_client)
    #example_stock_kline_query(data_client)
    #example_get_stock_dividend(data_client)
    #example_get_stock_lhb(data_client)
    #example_get_stock_lhb_by_date(data_client)
    #example_get_stock_dzjy(data_client)
    #example_get_stock_dzjy_by_date(data_client)
    #example_get_stock_bulletin(data_client)

